"""
RSI策略

基于相对强弱指标(RSI)的超买超卖策略。

作者: AI Assistant
版本: 1.0.0
日期: 2025-01-06
"""

import pandas as pd
import numpy as np
from backtest.strategy import StrategyBase


class RSIStrategy(StrategyBase):
    """
    RSI超买超卖策略
    
    使用RSI指标判断超买超卖状态：
    - RSI低于超卖线时买入
    - RSI高于超买线时卖出
    
    参数:
        period: RSI周期，默认14
        oversold: 超卖阈值，默认30
        overbought: 超买阈值，默认70
        
    信号规则:
        - RSI < 超卖阈值：买入
        - RSI > 超买阈值：卖出
        - 其他情况：持有
        
    示例:
        >>> strategy = RSIStrategy(period=14, oversold=30, overbought=70)
        >>> engine = BacktestEngine(strategy=strategy, data=data)
        >>> result = engine.run()
    """
    
    def __init__(
        self,
        period: int = 14,
        oversold: float = 30,
        overbought: float = 70
    ):
        """
        初始化RSI策略
        
        参数:
            period: RSI计算周期
            oversold: 超卖阈值（0-100）
            overbought: 超买阈值（0-100）
        """
        if period < 2:
            raise ValueError("RSI周期至少为2")
        
        if not 0 < oversold < overbought < 100:
            raise ValueError("需要满足: 0 < 超卖阈值 < 超买阈值 < 100")
        
        super().__init__(
            name=f"RSI策略(RSI{period}, {oversold}/{overbought})",
            min_period=period + 1,
            period=period,
            oversold=oversold,
            overbought=overbought
        )
        
        self.period = period
        self.oversold = oversold
        self.overbought = overbought
        self.last_signal = None
    
    def _calculate_rsi(self, data: pd.DataFrame) -> pd.Series:
        """
        计算RSI指标
        
        参数:
            data: 历史数据
            
        返回:
            pd.Series: RSI值序列
        """
        # 计算价格变化
        delta = data['收盘'].diff()
        
        # 分离上涨和下跌
        gain = (delta.where(delta > 0, 0)).rolling(
            window=self.period,
            min_periods=self.period
        ).mean()
        loss = (-delta.where(delta < 0, 0)).rolling(
            window=self.period,
            min_periods=self.period
        ).mean()
        
        # 计算RS和RSI
        rs = gain / loss
        rsi = 100 - (100 / (1 + rs))
        
        return rsi
    
    def generate_signal(self, data: pd.DataFrame, current_bar: pd.Series) -> str:
        """
        生成交易信号
        
        参数:
            data: 当前及之前的所有数据
            current_bar: 当前K线数据
            
        返回:
            str: 'BUY', 'SELL', 或 'HOLD'
        """
        # 确保有足够的数据
        if len(data) < self.period + 1:
            return 'HOLD'
        
        # 计算RSI
        rsi_series = self._calculate_rsi(data)
        
        # 获取当前RSI值
        if pd.isna(rsi_series.iloc[-1]):
            return 'HOLD'
        
        current_rsi = rsi_series.iloc[-1]
        
        # 超卖，买入信号
        if current_rsi < self.oversold:
            if self.last_signal != 'BUY':
                self.last_signal = 'BUY'
                return 'BUY'
        
        # 超买，卖出信号
        elif current_rsi > self.overbought:
            if self.last_signal != 'SELL':
                self.last_signal = 'SELL'
                return 'SELL'
        
        return 'HOLD'
    
    def get_indicator_values(self, data: pd.DataFrame) -> pd.DataFrame:
        """
        获取策略指标值（用于绘图）
        
        参数:
            data: 历史数据
            
        返回:
            pd.DataFrame: 包含RSI值的数据框
        """
        result = data.copy()
        result[f'RSI{self.period}'] = self._calculate_rsi(data)
        result['超卖线'] = self.oversold
        result['超买线'] = self.overbought
        return result

